// dqt(A)= (bit(A) == 1) ? 1 : -1
// C = dot(dqt(A), B)

template<
  int ts_cm, int ts_cn, int ts_ck,
  int wp_sz, int st_As_sz_per_t, int tk,
  int ld_Ag_wp_k, int ld_Ag_wp_m,
  int ld_Bg_wp_k, int ld_Bg_wp_n,
  int tk_ch_sz, int tm_sz, int tn_sz,
  int bm, int bn, int bt_sz, int wp_nb>
__global__ void mp_gemm_w1a8_sm75(
  const int m, const int n, const int k,
  const uint8_t * __restrict__ A, const int lda,
  const uint8_t * __restrict__ B, const int ldb,
  int32_t * __restrict__ C, const int ldc)
{
  const int bm_id = blockIdx.y;
  const int bn_id = blockIdx.x;
  const int tx = threadIdx.x;
  const int tk_sz = k/tk;
  const int bm_start = bm_id*bm;
  const int bn_start = bn_id*bn;
  const int tid = tx%32;
  const int wp_id = tx/32;
  const int wm_id = wp_id%tm_sz;
  const int wn_id = wp_id/tm_sz;
  const int ld_wm_sz = bm/ld_Ag_wp_m;
  const int ld_wn_sz = bn/ld_Bg_wp_n;
  const int ld_Ag_tm_id = tid/ld_Ag_wp_k;
  const int ld_Ag_tk_id = tid%ld_Ag_wp_k;
  const int ld_Bg_tn_id = tid/ld_Bg_wp_k;
  const int ld_Bg_tk_id = tid%ld_Bg_wp_k; 

  const int A_stmp_step = (tk/tk_ch_sz/st_As_sz_per_t);
  __shared__ uint32_t A_stmp[tk_ch_sz][bm*A_stmp_step];
  const int B_stmp_step = (tk/st_As_sz_per_t);
  __shared__ uint32_t B_stmp[bn*B_stmp_step];
  __shared__ int32_t C_stmp[bn*bm];

  A += bm_start*lda;
  B += bn_start*ldb;
  C += (bn_id*bn)*ldc+bm_start;
  uint8_t Ag_tmp = 0;
  int8_t As_tmp[8] = {0};

  wmma::fragment<wmma::matrix_a, ts_cm, ts_cn, ts_ck, int8_t, wmma::row_major> tsA;
  wmma::fragment<wmma::matrix_b, ts_cm, ts_cn, ts_ck, int8_t, wmma::col_major> tsB;
  wmma::fragment<wmma::accumulator, ts_cm, ts_cn, ts_ck, int32_t> tsC;
  wmma::fill_fragment(tsC, 0);
  __syncthreads();
  
  for(int tk_id= 0; tk_id < tk_sz; ++tk_id){ //block k loop
    const int tk_start = tk_id*tk;
    for(int wp_tid = wp_id; wp_tid < ld_wm_sz; wp_tid += wp_nb){//load A global
      const int bm_tid = wp_tid*ld_Ag_wp_m+ld_Ag_tm_id,
        k_id = tk_start+ld_Ag_tk_id*8;
      const int m_id = bm_start+bm_tid;
      if(m_id < m && k_id < k){
        Ag_tmp = (A + bm_tid*lda)[k_id/8];
      }else{
        Ag_tmp = 0;
      }
      // printf("Ag %x bm_id %d k_id %d \n", Ag_tmp, bm_tid, k_id);
      for(int8_t dq_id = 0; dq_id < 8; ++dq_id){
        As_tmp[dq_id] = (Ag_tmp&(1<<dq_id)) ? 1: -1;
      }
      for(int tk_cid = 0; tk_cid < tk_ch_sz; ++tk_cid){
        A_stmp[tk_cid][bm_tid*A_stmp_step+ld_Ag_tk_id] = 
          *(uint32_t*)(As_tmp+tk_cid*st_As_sz_per_t);
        // printf("As %x bm_id %d k_id %d \n", A_stmp[tk_cid][bm_tid*A_stmp_step+ld_Ag_tk_id], bm_tid, tk_start+tk_cid*16+ld_Ag_tk_id*4);
      }
    }
    for(int wp_tid = wp_id; wp_tid < ld_wn_sz; wp_tid += wp_nb){//load B global
      const int bn_tid = wp_tid*ld_Bg_wp_n+ld_Bg_tn_id,
        tk_id = ld_Bg_tk_id*st_As_sz_per_t, 
        k_id = tk_start+tk_id;
      const int n_id = bn_start+bn_tid;
      const int Bs_k_id = (((tk_id/4)%2)*16+((tk_id/8)*4))/4;
      if(n_id < n && k_id < k){
        B_stmp[bn_tid*B_stmp_step+Bs_k_id] = *(uint32_t*)(B+bn_tid*ldb+k_id);
      }else{
        B_stmp[bn_tid*B_stmp_step+Bs_k_id] = 0;
      }
      // printf("Bs %x bn_id %d k_id %d \n", B_stmp[bn_tid*B_stmp_step+Bs_k_id], bn_tid, tk_start+Bs_k_id*4);
    }
    __syncthreads();
    for(int tk_cid = 0; tk_cid < tk_ch_sz; ++tk_cid){
      wmma::load_matrix_sync(tsA, (int8_t*)(A_stmp[tk_cid]+(wm_id*ts_cm*A_stmp_step)), 16);
      wmma::load_matrix_sync(tsB, (int8_t*)(B_stmp+(wn_id*ts_cn*B_stmp_step)+tk_cid*4), 32);
      wmma::mma_sync(tsC, tsA, tsB, tsC);
    }
    __syncthreads();
  }
  wmma::store_matrix_sync((int32_t*)(C_stmp+(wn_id*ts_cn)*bm+wm_id*ts_cm), tsC, bm, wmma::mem_col_major); //store C shared
  // if(wn_id == 0 && wm_id == 0){
  //   for(int t=0; t<tsC.num_elements; t++){
  //     printf("t %d c %d \n", t, tsC.x[t]);
  //   }
  // }
  __syncthreads();
  for(int tid = tx; tid < bn*bm; tid += bt_sz){ //store C global
    const int bn_tid =  tid/bm, bm_tid = tid%bm;
    const int n_id = bn_start+bn_tid, m_id = bm_start+bm_tid;
    if(n_id < n && m_id < m){
      C[bn_tid*ldc+bm_tid] = C_stmp[bn_tid*bm+bm_tid];
    }
    // printf("n_id %d m_id %d c_stmp %d \n", n_id, m_id, C_stmp[bn_tid*bm+bm_tid]);
  }
}

/**
 * Weight: A (int1)
 * Input Activation: B (int8)
 * Output Activation: C (int32)
*/
cublasStatus_t mp_gemm_w1a8(
  const int m, const int n, const int k,
  void *A, const int lda, 
  void *B, const int ldb, 
  void *C, const int ldc,
  cudaStream_t cuStream = nullptr)
{
  using AT = uint8_t;
  using BT = uint8_t;
  using CT = int32_t;
  BT *d_B = (BT*)B;
  CT *d_C = (CT*)C;
  AT *d_A = (AT*)A;

  const int ts_cm = 32, ts_cn = 8, ts_ck = 16;
  const int wp_sz = 32;
  const int st_As_sz_per_t = 4, tk = 32;
  const int ld_Ag_wp_k = tk/8, ld_Ag_wp_m = wp_sz/ld_Ag_wp_k;
  const int ld_Bg_wp_k = tk/4, ld_Bg_wp_n = wp_sz/ld_Bg_wp_k;
  const int tk_ch_sz = 8/st_As_sz_per_t;
  const int tm_sz = 2, tn_sz = 2;
  const int bm = tm_sz*ts_cm, bn = tn_sz*ts_cn;
  const int bt_sz = wp_sz*tm_sz*tn_sz;
  const int wp_nb = tm_sz*tn_sz;

  dim3 g_sz(DIV_UP(n, bn), DIV_UP(m, bm)), b_sz(bt_sz);
  mp_gemm_w1a8_sm75<
    ts_cm, ts_cn, ts_ck, wp_sz, st_As_sz_per_t, tk, ld_Ag_wp_k, ld_Ag_wp_m, 
      ld_Bg_wp_k, ld_Bg_wp_n, tk_ch_sz, tm_sz, tn_sz, bm, bn, bt_sz, wp_nb>
    <<<g_sz, b_sz, 0, cuStream>>>
    (m, n, k, d_A, lda, d_B, ldb, d_C, ldc);
  {
    static int num = 0;
    if(num == 0){
      printf("=================>  \n");
      printf(" thread sz %d/2560 \n", g_sz.x*g_sz.y*b_sz.x);
      printf(" st_As_sz_per_t %d tk %d ld_Ag_wp_k %d ld_Ag_wp_m %d ld_Bg_wp_k %d \n ld_Bg_wp_n %d tk_ch_sz %d tm_sz %d tn_sz %d \n bm %d bn %d bt_sz %d wp_nb %d \n", 
        st_As_sz_per_t, tk, ld_Ag_wp_k, ld_Ag_wp_m, ld_Bg_wp_k, ld_Bg_wp_n, tk_ch_sz, tm_sz, tn_sz, bm, bn, bt_sz, wp_nb);
      printf("=================>  \n");
    }
    ++num;
  }
  return CUBLAS_STATUS_SUCCESS;
}